"""Coordinate transform and quaternion/euler helpers extracted from robot_client.py.

Provides:
 - rotation_matrix_to_rpy_quat_rad_and_xyzw
 - quat_to_rot_matrix
 - transform_pose_camera_to_ros
 - euler_rad_to_deg

These functions are pure-numeric and have no ROS/ZMQ side-effects.
"""
from __future__ import annotations
import numpy as np

def rotation_matrix_to_rpy_quat_rad_and_xyzw(rot):
    """From a 3x3 rotation matrix return (roll,pitch,yaw) in radians and quaternion in [x,y,z,w].

    Returns (rpy_rad_tuple, quat_xyzw_tuple) or (None, None) on failure.
    """
    try:
        R = np.array(rot, dtype=float)
        r11, r12, r13 = R[0, 0], R[0, 1], R[0, 2]
        r21, r22, r23 = R[1, 0], R[1, 1], R[1, 2]
        r31, r32, r33 = R[2, 0], R[2, 1], R[2, 2]

        # roll (x)
        roll = np.arctan2(r32, r33)
        # pitch (y)
        pitch = np.arctan2(-r31, np.sqrt(r32 * r32 + r33 * r33))
        # yaw (z)
        yaw = np.arctan2(r21, r11)

        # compute quaternion (w, x, y, z) using same algorithm
        tr = r11 + r22 + r33
        if tr > 0:
            S = np.sqrt(tr + 1.0) * 2.0
            qw = 0.25 * S
            qx = (r32 - r23) / S
            qy = (r13 - r31) / S
            qz = (r21 - r12) / S
        else:
            if (r11 > r22) and (r11 > r33):
                S = np.sqrt(1.0 + r11 - r22 - r33) * 2.0
                qw = (r32 - r23) / S
                qx = 0.25 * S
                qy = (r12 + r21) / S
                qz = (r13 + r31) / S
            elif r22 > r33:
                S = np.sqrt(1.0 + r22 - r11 - r33) * 2.0
                qw = (r13 - r31) / S
                qx = (r12 + r21) / S
                qy = 0.25 * S
                qz = (r23 + r32) / S
            else:
                S = np.sqrt(1.0 + r33 - r11 - r22) * 2.0
                qw = (r21 - r12) / S
                qx = (r13 + r31) / S
                qy = (r23 + r32) / S
                qz = 0.25 * S

        # return euler in radians and quaternion as [x,y,z,w]
        quat_xyzw = (float(qx), float(qy), float(qz), float(qw))
        rpy_rad = (roll, pitch, yaw)
        return rpy_rad, quat_xyzw
    except Exception:
        return None, None


def quat_to_rot_matrix(q):
    """Convert quaternion to 3x3 rotation matrix.

    Accepts q as iterable of 4 numbers in either [x,y,z,w] or [w,x,y,z] order.
    Heuristic: if the last element has the largest absolute value we assume [x,y,z,w],
    otherwise assume [w,x,y,z]. Returns 3x3 numpy array.
    """
    q = np.array(q, dtype=float)
    if q.size != 4:
        raise ValueError('Quaternion must have 4 elements')

    # detect ordering
    if abs(q[3]) >= abs(q[0]) and abs(q[3]) >= abs(q[1]) and abs(q[3]) >= abs(q[2]):
        x, y, z, w = q[0], q[1], q[2], q[3]
    else:
        w, x, y, z = q[0], q[1], q[2], q[3]

    # normalize
    n = np.linalg.norm([w, x, y, z])
    if n == 0:
        raise ValueError('Zero-norm quaternion')
    w, x, y, z = w / n, x / n, y / n, z / n

    R = np.array([
        [1 - 2 * (y * y + z * z),     2 * (x * y - z * w),     2 * (x * z + y * w)],
        [    2 * (x * y + z * w), 1 - 2 * (x * x + z * z),     2 * (y * z - x * w)],
        [    2 * (x * z - y * w),     2 * (y * z + x * w), 1 - 2 * (x * x + y * y)]
    ], dtype=float)
    return R


def transform_pose_camera_to_ros(position, rotation_matrix=None, quaternion=None):
    """Transform pose from camera/OpenCV frame (x:right, y:down, z:forward)
    to ROS-style frame (x:forward, y:left, z:up).

    Returns: (position_new, quat_xyzw, euler_xyz_rad)
    """
    # orinal
    T = np.array([
        [0, 0, 1],
        [-1, 0, 0],
        [0, -1, 0]
    ], dtype=float)

    pos = np.array(position, dtype=float).reshape(3)
    pos_new = T @ pos

    if rotation_matrix is None:
        if quaternion is None:
            raise ValueError('Either rotation_matrix or quaternion must be provided')
        rot = quat_to_rot_matrix(quaternion)
    else:
        rot = np.array(rotation_matrix, dtype=float).reshape((3, 3))

    rot_new = T @ rot 

    euler_rad, quat_xyzw = rotation_matrix_to_rpy_quat_rad_and_xyzw(rot_new)
    q = np.array(quat_xyzw, dtype=float)
    qn = np.linalg.norm(q)
    if qn == 0:
        raise ValueError('Resulting quaternion has zero norm')
    q = (q / qn).tolist()

    return pos_new, (q[0], q[1], q[2], q[3]), euler_rad


def euler_rad_to_deg(euler_rad):
    """Convert an (roll,pitch,yaw) tuple in radians to degrees.

    Returns None if input is None.
    """
    if euler_rad is None:
        return None
    return tuple(np.degrees(euler_rad))
